/**
 * @file offb_node.cpp
 * @brief Offboard控制节点
 * 起飞后寻找目标，然后飞向目标并投弹，最后降落
 * 杨朋飞 2023.5.9
 */

#include <ros/ros.h>
#include <std_msgs/String.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/TwistStamped.h>
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>
#include <mavros_msgs/AttitudeTarget.h>
#include <tf/transform_datatypes.h>
#include <tf/tf.h>
#include <tf/transform_datatypes.h>

//mavros/state topic callback
mavros_msgs::State current_state;
void state_cb(const mavros_msgs::State::ConstPtr& msg){
    current_state = *msg;
}

//local_position/pose topic callback
geometry_msgs::PoseStamped local_position;
double roll_actual,pitch_actual,yaw_actual;
void local_position_cb(const geometry_msgs::PoseStamped::ConstPtr& msg)
{
    local_position = *msg;
    tf::Quaternion quat;
    //以下部分为四元数转欧拉角
    tf::quaternionMsgToTF(local_position.pose.orientation, quat);
    tf::Matrix3x3(quat).getRPY(roll_actual,pitch_actual,yaw_actual);
    //需进行以下变换，才与地面站上结果一致：
    yaw_actual = 90 - (yaw_actual * 180 / M_PI);
    //将偏航角范围调整至[0,360]
    if (yaw_actual >= 360.0) yaw_actual -= 360.0;
    if (yaw_actual < 0.0) yaw_actual += 360.0;
}

//talker topic callback
geometry_msgs::PoseStamped opencv_data;
void talker_cb(const geometry_msgs::PoseStamped::ConstPtr& msg)
{
    opencv_data = *msg;
}

//bomber topic callback
bool bombing_finished = false;
ros::Time bombing_finish_time = ros::Time(0);
void bomber_cb(const std_msgs::String::ConstPtr& msg)
{
    if (msg->data == "bomber_finish")
    {
        if (bombing_finished == false)
            bombing_finish_time = ros::Time::now();
        bombing_finished = true;
        ROS_INFO("Bomb dropped");
    }
}


int main(int argc, char **argv)
{
    /*-----ROS Subscribers & Publishers-----*/
    ros::init(argc, argv, "offb_node");
    ros::NodeHandle nh;

    ros::Subscriber state_sub = nh.subscribe<mavros_msgs::State>
            ("mavros/state", 10, state_cb);
    ros::Publisher local_pos_pub = nh.advertise<geometry_msgs::PoseStamped>
            ("mavros/setpoint_position/local", 10);
    ros::Publisher setpoint_velocity_pub = nh.advertise<geometry_msgs::TwistStamped>
            ("mavros/setpoint_velocity/cmd_vel", 10);
    ros::Publisher start_bombing_pub = nh.advertise<std_msgs::String>
            ("bomber/start", 10);
    ros::Subscriber bombing_finish_sub = nh.subscribe<std_msgs::String>
            ("bomber/finish", 5, bomber_cb);
    ros::ServiceClient arming_client = nh.serviceClient<mavros_msgs::CommandBool>
            ("mavros/cmd/arming");
    ros::ServiceClient set_mode_client = nh.serviceClient<mavros_msgs::SetMode>
            ("mavros/set_mode");
    ros::Subscriber local_pos_sub = nh.subscribe<geometry_msgs::PoseStamped>
            ("mavros/local_position/pose", 10, local_position_cb);
    ros::Subscriber talker_sub = nh.subscribe<geometry_msgs::PoseStamped>
            ("talker", 10, talker_cb);            

    //the setpoint publishing rate MUST be faster than 2Hz
    ros::Rate rate(20.0);

    /*-----Preflight-----*/
    // wait for FCU connection
    while(ros::ok() && !current_state.connected){
        ros::spinOnce();
        rate.sleep();
    }

    const double DEG2RAD = 0.01745329252;
    geometry_msgs::PoseStamped pose = {};
    geometry_msgs::TwistStamped setpoint_vel = {};
    double yaw_setpoint = 0;
    pose.pose.orientation.w = cos((90-yaw_setpoint)*DEG2RAD/2);
    pose.pose.orientation.z = sin((90-yaw_setpoint)*DEG2RAD/2);
    double initial_heading = 0;
    double initial_x = 0;
    double initial_y = 0;
    double initial_z = 0;
    std::stringstream ss;
    std_msgs::String bomber_msg;
    ss.str("stop");
    bomber_msg.data = ss.str();

    //send a few setpoints before starting
    for(int i = 20; ros::ok() && i > 0; --i){
        local_pos_pub.publish(pose);
        ros::spinOnce();
        rate.sleep();
    }

    for(int i = 20; ros::ok() && i > 0; --i){
        local_pos_pub.publish(pose);
        start_bombing_pub.publish(bomber_msg);
        ros::spinOnce();
        rate.sleep();
        initial_heading += yaw_actual;
        initial_x += local_position.pose.position.x;
        initial_y += local_position.pose.position.y;
        initial_z += local_position.pose.position.z;
    }
    
    //Change setpoint to the average of initial position & heading
    initial_heading = initial_heading / 20;
    initial_x = initial_x / 20;
    initial_y = initial_y / 20;
    initial_z = initial_z / 20;
    yaw_setpoint = initial_heading;
    pose.pose.orientation.w = cos((90-yaw_setpoint)*DEG2RAD/2);
    pose.pose.orientation.z = sin((90-yaw_setpoint)*DEG2RAD/2);
    pose.pose.position.x = initial_x;
    pose.pose.position.y = initial_y;
    pose.pose.position.z = initial_z + 5;
    ROS_INFO("Initial position x: %.1lf, y: %.1lf, z:%.1lf, heading: %.1lf", initial_x, initial_y, initial_z, initial_heading);
    ROS_INFO("Setpoint position x: %.1lf, y: %.1lf, z:%.1lf, heading: %.1lf", pose.pose.position.x, pose.pose.position.y, pose.pose.position.z, initial_heading);

    /*-----offboard & arm-----*/
    mavros_msgs::SetMode offb_set_mode;
    offb_set_mode.request.custom_mode = "OFFBOARD";
    mavros_msgs::CommandBool arm_cmd;
    arm_cmd.request.value = true;
    ros::Time last_request = ros::Time::now();
    while(ros::ok()){//只会执行一次，避免起飞后抢权限
        if( current_state.mode != "OFFBOARD" &&
            (ros::Time::now() - last_request > ros::Duration(4.0))){
            if( set_mode_client.call(offb_set_mode) &&
                offb_set_mode.response.mode_sent){
                ROS_INFO("Offboard enabled");
                last_request = ros::Time::now();
                break;
            }
            last_request = ros::Time::now();
        }
        local_pos_pub.publish(pose);
 
        ros::spinOnce();
        rate.sleep();
    }

    while(ros::ok() && !current_state.armed){
        if(ros::Time::now() - last_request > ros::Duration(4.0))
        {
            if( arming_client.call(arm_cmd) &&
                arm_cmd.response.success){
                ROS_INFO("Arm vehicle success");
                last_request = ros::Time::now();
                break;
            }
            last_request = ros::Time::now();
        }
        local_pos_pub.publish(pose);
        
        ros::spinOnce();
        rate.sleep();
    }
    ROS_INFO("Vehicle armed");

    /*-----找到目标并锁定-----*/
    ros::Time last_found_target = ros::Time::now();
    ros::Time locked_time = ros::Time(0);
    while(ros::ok())
    {
        if (fabs(local_position.pose.position.z - pose.pose.position.z) < 0.1)
        {
            if (opencv_data.pose.position.x > 0 && opencv_data.pose.position.y > 0)
            {
                yaw_setpoint = yaw_actual + (opencv_data.pose.position.x - 328) * 15 / 328;
                if (yaw_setpoint >= 360.0) yaw_setpoint -= 360.0;
                if (yaw_setpoint < 0.0) yaw_setpoint += 360.0;
                pose.pose.orientation.w = cos((90-yaw_setpoint)*DEG2RAD/2);
                pose.pose.orientation.z = sin((90-yaw_setpoint)*DEG2RAD/2);
                last_found_target = ros::Time::now();
                if ((yaw_setpoint > 5.0 && yaw_setpoint < 355.0 && fabs(yaw_setpoint - yaw_actual) < 5.0)
                    || (yaw_setpoint <= 5.0 && ((355.0 + yaw_setpoint < yaw_actual) || yaw_actual < yaw_setpoint + 5.0))
                    || (yaw_setpoint >= 355.0 && ( yaw_actual > yaw_setpoint - 5.0 || yaw_actual < yaw_setpoint - 355.0)))//setpoint与actual夹角小于5度
                {
                    if (locked_time == ros::Time(0))
                        locked_time = ros::Time::now();
                    else
                    {                        
                        if (ros::Time::now() - locked_time > ros::Duration(3.0))
                        {
                            ROS_INFO("Locked on target!");
                            break;
                        }
                    }
                }
                else
                    locked_time = ros::Time(0);
            }
            else if (ros::Time::now() - last_found_target > ros::Duration(3.0))
            {
                yaw_setpoint += 1;
                if (yaw_setpoint >= 360.0) yaw_setpoint -= 360.0;
                if (yaw_setpoint < 0.0) yaw_setpoint += 360.0;
                pose.pose.orientation.w = cos((90-yaw_setpoint)*DEG2RAD/2);
                pose.pose.orientation.z = sin((90-yaw_setpoint)*DEG2RAD/2);
                locked_time = ros::Time(0);
            }
            else
                locked_time = ros::Time(0);
        }

        local_pos_pub.publish(pose);
        if(!(current_state.mode == "OFFBOARD" || current_state.mode == ""))
        {
            ROS_ERROR("Vehicle now in %s mode, quitting.", current_state.mode.c_str());
            std::cout << "Mode: " << current_state.mode << std::endl;
            return -1;
        }
        ros::spinOnce();
        rate.sleep();
    }

    /*-----改为速度控制并飞向目标-----*/
    double abs_set_vel = 1.5;//Set ground speed 1.5m/s
    setpoint_vel.twist.linear.z = 0;
    double setpoint_ang_vel_z = 0;
    ss.str("start");
    bomber_msg.data = ss.str();
    bombing_finish_time = ros::Time(0);
    bombing_finished = false;
    while(ros::ok())
    {
        if (opencv_data.pose.position.x > 0 && opencv_data.pose.position.y > 0)
        {
            yaw_setpoint = yaw_actual + (opencv_data.pose.position.x - 328) * 15 / 328;
            if (yaw_setpoint >= 360.0) yaw_setpoint -= 360.0;
            if (yaw_setpoint < 0.0) yaw_setpoint += 360.0;
            last_found_target = ros::Time::now();
        }
        else if (!bombing_finished && ros::Time::now() - last_found_target > ros::Duration(5.0))//找不到目标，而又没有收到投弹完成消息
        {
            ROS_WARN("Lost target, landing now!");
            break;
        }
        if (bombing_finished && ros::Time::now() - bombing_finish_time > ros::Duration(1.5))//投弹完成后1.5s返航
        {
            ROS_INFO("Bombing completed, landing now!");
            break;
        }
        setpoint_vel.twist.linear.x = abs_set_vel * sin(yaw_setpoint * DEG2RAD);
        setpoint_vel.twist.linear.y = abs_set_vel * cos(yaw_setpoint * DEG2RAD);
        setpoint_ang_vel_z = yaw_actual - yaw_setpoint;
        if (setpoint_ang_vel_z < -180.0)
            setpoint_ang_vel_z += 360.0;
        else if (setpoint_ang_vel_z > 180.0)
            setpoint_ang_vel_z -= 360.0;
        setpoint_vel.twist.angular.z = setpoint_ang_vel_z * DEG2RAD * 3;
        setpoint_velocity_pub.publish(setpoint_vel);
        start_bombing_pub.publish(bomber_msg);
        if(!(current_state.mode == "OFFBOARD" || current_state.mode == ""))
        {
            ROS_ERROR("Vehicle now in %s mode, quitting.", current_state.mode.c_str());
            std::cout << "Mode: " << current_state.mode << std::endl;
            return -1;
        }
        ros::spinOnce();
        rate.sleep();
    }

    /*-----完成任务，飞回起飞点降落-----*/
    yaw_setpoint = initial_heading;
    pose.pose.orientation.w = cos((90-yaw_setpoint)*DEG2RAD/2);
    pose.pose.orientation.z = sin((90-yaw_setpoint)*DEG2RAD/2);
    pose.pose.position.x = initial_x;
    pose.pose.position.y = initial_y;
    pose.pose.position.z = initial_z + 5;
    while(ros::ok())
    {
        if (sqrt((local_position.pose.position.x - pose.pose.position.x) * (local_position.pose.position.x - pose.pose.position.x)
         + (local_position.pose.position.y - pose.pose.position.y) * (local_position.pose.position.y - pose.pose.position.y)
         + (local_position.pose.position.z - pose.pose.position.z) * (local_position.pose.position.z - pose.pose.position.z)) < 0.15)
        {
            offb_set_mode.request.custom_mode = "AUTO.LAND";
            if (current_state.mode != "AUTO.LAND" && (ros::Time::now() - last_request > ros::Duration(5.0)))
            {

                if (set_mode_client.call(offb_set_mode) && offb_set_mode.response.mode_sent)
                {
                    ROS_INFO("AUTO.LAND enabled");
                }
                last_request = ros::Time::now();
            }
        }
        else
        {
            local_pos_pub.publish(pose);
        }
        if(!(current_state.mode == "OFFBOARD" || current_state.mode == "AUTO.LAND" || current_state.mode == ""))
        {
            ROS_ERROR("Vehicle now in %s mode, quitting.", current_state.mode.c_str());
            return -1;
        }
        ros::spinOnce();
        rate.sleep();
    }

    ROS_INFO("Done!");

    return 0;
}


